import streamlit
import streamlit as st

print("start refresh")

if "stage" not in st.session_state:
    st.session_state.stage = 1

# 错误的示例
# if st.session_state.stage == 1:
#     st.write("Stage 1")
#     if st.button("Click to Stage 2"):
#         print("click button 1")
#         st.session_state.stage = 2
# if st.session_state.stage == 2:
#     st.write("Stage 2")
#     if st.button("Click to Stage 1"):
#         print("click button 2")
#         st.session_state.stage = 1
# else:
#     print("not found")
#正确的示例
def goto_step(stage):
    print("click button:",(2-stage+1))
    st.session_state.stage = stage

if st.session_state.stage == 1:
    st.write("Stage 1")
    st.button("Click to Stage 2",on_click=goto_step,args=(2,))

if st.session_state.stage == 2:
    st.write("Stage 2")
    st.button("Click to Stage 1",on_click=goto_step,args=(1,))

st.write(st.session_state)
